Skip to content

V1.0 - #752

Open
thomaspinder wants to merge 111 commits into
mainfrom
v1.0
Open

V1.0#752
thomaspinder wants to merge 111 commits into
mainfrom
v1.0

Conversation

@thomaspinder

Copy link
Copy Markdown
Owner

Checklist

  • I've formatted the new code by running uv run poe format before committing.
  • I've added tests for new code.
  • I've added docstrings for the new code.

Description

Please describe your changes here. If this fixes a bug, please link to the issue, if possible.

Issue Number: N/A

thomaspinder and others added 30 commits August 6, 2026 01:27
conjugate_mll had no value-level test anywhere in the suite, yet serves as
the oracle for the Kalman MLL and collapsed_elbo. These closed-form pins,
computed through an independent jnp.linalg path, give the reference frame
its ground truth ahead of the v1.0 conditioning refactor.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
…tter bug

collapsed_elbo(z=X) vs conjugate_mll and whitened-vs-unwhitened predicts at
matched parameters now guard the five independent derivations of the
conjugate conditioning algebra. The strict xfail documents that at
non-default jitter the derivations factorise different matrices — the bug
the v1.0 conditioning module removes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
_compare previously swallowed AssertionError with a print, so the harness
could never fail. Failures are now collected per-example and raised at the
end of test(), making the four golden-value pins a real no-behaviour-change
net for the v1.0 refactor.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
The newly-loud harness exposed pre-existing drift in all four examples: the
collapsed/uncollapsed goldens predated the real-data example swap (#696) and
the regression/heteroscedastic goldens predated subsequent behaviour fixes
(#707/#708/#713 and dependency bumps). The toothless harness never noticed.
Re-pinned so the net measures the v1.0 refactor, not history.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
The full-dataset size a minibatch ELBO needs now travels on the one object
that knows it, as static pytree aux_data, instead of being smuggled through
likelihood constructors.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
…nd noise_prior

Of ~245 occurrences of num_datapoints, only four were real reads: the ELBO
minibatch scale (now served by Dataset.n_total) and latent sizing (moves to
data-contact time in the JointModel rewrite). No likelihood used the value
internally, and nothing validated it — a wrong value silently mis-scaled
the ELBO. noise_prior moves to the model layer, where priors live.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
The API now mirrors the maths: prior * likelihood -> JointModel (the joint
p(f,y), the trainable object); model.condition(D) — sugar: model | D —
returns an immutable Posterior pytree caching the Cholesky factor and
representer weights. The predictive, log_marginal_likelihood, loo, and
pathwise sample_approx are views of that one factorisation, deleting the
eleven independent derivations and the two-owner jitter split
(prior.jitter is now the single knob, applied once inside conditioning).

- gpjax/conditioning.py: deep module (Posterior, ExactPosterior,
  LatentPosterior); MO validation moves to condition time; sample_approx
  refuses multi-output loudly instead of silently broadcasting wrong.
- gps.py: Prior (AbstractPrior folded in), ConjugateModel,
  NonConjugateModel (lazy latent, sized at data contact),
  HeteroscedasticModel (owns noise_prior — likelihoods are pure
  conditionals again, killing the likelihoods->gps circular import).
  Deleted: AbstractPrior, AbstractPosterior, LatentPosterior marker,
  ChainedPosterior marker, construct_posterior (now construct_model).
- objectives: conjugate_mll/conjugate_loocv/log_posterior_density are
  one-line views of the conditioned posterior.
- fit: _prepare_model hook sizes lazily-initialised state from data.
- predict(t, D) survives as documented one-line sugar everywhere.
- return_covariance_type kwarg renamed to covariance.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
Mechanical: ConjugatePosterior->ConjugateModel and friends,
construct_posterior->construct_model, return_covariance_type->covariance,
num_datapoints/noise_prior constructor ceremony deleted (~240 sites).
Semantic: heteroscedastic tests build HeteroscedasticModel directly;
non-conjugate tests size the latent via init_latent; docs/index.md
quickstart shows condition(); regression example narrates the
condition API; StateSpaceConjugatePosterior renamed
StateSpaceConjugateModel.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
…ault jitter

The model-side two-owner jitter bug is fixed; the strict xfail narrows to
the family-side knob, which unifies in the variational stack PR.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
- reference/gps.md lists the JointModel hierarchy and the conditioning
  module; state_space.md and linalg.md updated for renamed/new symbols
- stale glossary/sharp_bits/classification xrefs renamed
- poisson example initialises the lazy non-conjugate latent before MCMC
- ADR directory excluded from the docs site (in-repo records for now)
- codeautolink match_block warnings suppressed on every path: a matcher
  limitation on doctest-SKIP blocks, predating this stack — the docs
  workflow had not run cold since the Sphinx migration, so tonight's PR
  pushes surfaced it for the first time

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
…inery

Implements Salimbeni et al. 2018 (arXiv:1803.09151) natural-gradient VI for
VariationalGaussian and WhitenedVariationalGaussian.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
Adversarial review of the natural-gradient core turned up two correctness
defects, one performance defect, three contract mismatches and a set of
documentation and test gaps. Fixes, in order of severity:

* `_first_valid_trial` leaked float64 into the scan carry. Under x64 the
  exponent `jnp.arange(K + 1)` is int64, so `backoff ** arange` was a non-weak
  float64 that promoted a float32 model and made `lax.scan` reject the carry.
  The trial ladder is now cast to the dtype of Theta_2, so `fit_natgrads` is no
  longer strictly narrower than `fit`.

* The backoff replicated the whole theta -> xi map across all K+1 trials, at a
  measured 13% of total training wall clock at M=200 -- not the "negligible"
  cost impl-plan 1.2.5 assumed. Only the admissibility probe is replicated now;
  the inversion, the X^T X product and the second Cholesky run once, at the
  accepted step size. Measured overhead at M=200 falls to 4.6%.

* `fit_natgrads`' signature rejected everything `_check_natgrad_lr` blessed:
  `natgrad_lr=1`, `map_jitter=0`, `backoff=1` and a 0-d array all raised under
  the beartype import hook. Annotations widened, the validator now accepts 0-d
  arrays and rejects bool, and the entry point (not just the validator) is
  tested with each.

* `_reject_frozen_coordinates` matched coordinates against top-level dataclass
  fields by identity, so a future nested registration would have passed the
  guard silently. It now re-walks the tree with the selector as the `is_leaf`
  predicate and reports the full key path. Its message also pluralises and
  points at a remedy that exists -- the old one recommended freezing the whole
  family, which re-raises the same error.

* `fit_natgrads` now calls the guard under `safe=True` as impl-plan 1.2.4
  prescribes, and forwards `log_rate` to `vscan` instead of documenting a knob
  that did nothing.

Tests: `test_natgrad_backoff_recovers_from_large_step` never exercised the
backoff (the k=0 trial was already admissible at natgrad_lr=100), so it now
starts from a tight S_0, asserts the un-shrunk step genuinely leaves the cone,
and checks the accepted step size. The Cholesky-budget test could not see the
vmap width; it is now an exact count parametrised over max_backoff, plus a
lowered-IR test asserting the batched Cholesky has leading dimension K+1. Added
dtype-preservation tests, and moved the duplicated conjugate oracle into
`tests/_reference/conjugate_svgp.py` so the two transcriptions cannot drift.

Docs: all seven exported functions gained runnable `Example:` blocks (impl-plan
section 6), the map_jitter bias on `history` and the eta -> xi cancellation
regime are now documented on the public surface, and two docstrings became raw
strings.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
…_natgrads

The ```pycon fences render fine under MkDocs but are invalid RST under
Sphinx/napoleon, producing docutils warnings that fail the -W docs-ci
gate. Drop the fences; the indented doctest block matches fit()'s style.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
Adds examples/dual_svgp.py, the second of the two natural-gradient tutorial
notebooks, and wires both of them into the documentation.

The notebook derives the dual parameterisation for a reader who has been
through examples/natgrads.py: the additive split eta = eta_0(theta) + lambda,
the EP-style likelihood sites and their tying to inducing space, the two
convention traps (flanked vs un-flanked storage, and the -1/2 on Lambda_2),
and the tied natural-gradient update, which is an affine convex combination on
the stored sites because grad_mu KL == lambda exactly, so the KL is never
differentiated and no theta <-> eta round trip is needed.

Measured in the executed run:

* one rho = 1 full-batch step on a conjugate model with a non-zero mean
  function reproduces the Titsias optimum to 1.7e-12 (mean) and 1.2e-13
  (covariance), a second step moves nothing, and dual_elbo matches the
  collapsed bound up to exactly N * jitter / (2 sigma^2);
* rho is gamma: matched dual and Salimbeni E-steps agree in (m, S) to 3.1e-15
  over six full-batch steps at rho in {0.3, 0.8, 1.0}, and two frozen-
  hyperparameter fit_natgrads runs overlay to 4.3e-14 over 50 iterations;
* the banana benchmark (N = 2000, M = 50, B = 256, 1000 iterations, the same
  make_banana and jr.key(42) as the natural-gradients notebook) with t-SVGP,
  the Salimbeni natgrad and Adam alone, per iteration and per wall-clock
  second -- both natural-gradient runs reach Adam's 1000-iteration bound at
  iteration 109, and the dual iteration is not cheaper at this scale;
* hyperparameter learning: the dual/standard hyper-gradient gap falls from
  3.9e1 to 8.5e-15 as the E-step converges; dual dominance is not uniform when
  sparse (negative gaps at M = 5 and M = 10) but holds by 1.5e5 nats at Z = X;
  and a 40-round VEM loop ends 0.50 nats ahead on dual_elbo with equal
  held-out NLPD.

Docs wiring: both notebooks added to the mkdocs.yml Tutorials nav and to
CTA_NOTEBOOKS in docs/scripts/gen_examples.py, plus an adam2021dual entry in
docs/refs.bib.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
Corrections to examples/dual_svgp.py, all re-verified against a fresh
end-to-end execution:

* the intro no longer implies the two hyperparameter gradients agree at
  theta_t; they agree only at a converged E-step, which is what the
  notebook's own gradient table measures;
* added a notation-reconciliation note bridging the natural-gradients
  notebook's (theta, eta, lambda) to this one's (eta, mu, lambda), and
  restated the borrowed identity and H_2 in these letters;
* the c(theta) remark now names the site convention it holds under
  (normalised projected) and fixes the sign apposition;
* the bound-slice prose is now asymmetric, as the data are: l collapses
  on the long-lengthscale side while l-bar barely moves, but both
  collapse together on the short side, where the sparse approximation
  itself has failed. Added edge diagnostics to back it;
* the banana benchmark now states which run ends ahead and bounds what
  that comparison can mean;
* the VEM panel plots the round-by-round bound lead rather than two
  indistinguishable traces. That exposed a false claim: the dual M-step
  is behind for the first seven rounds, crosses at round 8 and holds a
  sub-nat lead thereafter. Prose corrected and the crossing printed;
* order-of-magnitude claims restated from the printed values (Titsias
  agreement, the jitter residual now printed to twelve digits, the
  banana condition-number ratio, the M = 20 crossing-point noise);
* the dominance row now names conjugacy as well as Z = X;
* the roadmap names the two sections it had omitted;
* the banana-copy rationale no longer overstates cross-notebook
  comparability, and the wall-clock explanation leads with the
  O(M^3)-vs-O(BM^2) argument rather than an evaluation XLA folds away;
* ruff-format clean, no code line over 88 characters.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
…tionVariationalGaussian

Both classes were parameterisation-only: they stored the natural or
expectation coordinates of q(u) but shipped no way to take a
natural-gradient step in them, so they bought nothing over the standard
families.

Natural-gradient geometry belongs to the optimiser, not the family. The
Fisher matrix is exactly the Jacobian dn/dt, so the natural gradient with
respect to the natural parameters equals the ordinary gradient with
respect to the expectation parameters, in any parameterisation. fit_natgrads
(PR#1, gpjax/natural_gradients.py) therefore computes the transforms
on the fly and operates directly on VariationalGaussian and
WhitenedVariationalGaussian, which store constraint-respecting
coordinates. Users of the removed classes should switch to
VariationalGaussian with gpjax.fit_natgrads.

Also drops the now-dead _psd helper and the cholesky_factor import, whose
only call sites lived inside the deleted classes, and the "natural" and
"expectation" arms of the VariationalParametrisationSuite ASV benchmark.

BREAKING CHANGE: NaturalVariationalGaussian and ExpectationVariationalGaussian
are removed from gpjax.variational_families.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
Implements the dual/site parameterisation of sparse variational GPs from
Adam, Chang, Khan and Solin (2021), "Dual Parameterization of Sparse
Variational Gaussian Processes", NeurIPS 2021 (arXiv:2111.03412).

`DualVariationalGaussian` stores an unnormalised Gaussian site on the
*centred* inducing outputs -- `dual_vector` is the site's first natural
parameter and `dual_matrix` its precision, both `Real` and both defaulting
to zero, so q(u) = p(u) at initialisation. Neither carries a constraining
bijection: PSD-ness of the site precision comes from the convex-combination
structure of the natural-gradient update, and a bijection would destroy that
affine step. Moments, marginals, prior KL and predictions all route through
the working matrix R = Kzz + Kzz L2 Kzz, which dominates Kzz and is therefore
always factorisable even when the site precision is rank deficient; exactly
two Cholesky factorisations are taken per call and nothing is inverted.

The centred convention is the correction to the reference implementation's
mean-function bug, which shifts by `predict_f(Z)` and so is wrong for any
non-zero mean function. `test_dual_natgrad_handles_non_zero_mean_function`
pins the correct behaviour and asserts that the uncentred variant misses.

`marginals` adds the family's jitter to every marginal variance. This is
load-bearing, not cosmetic: `VariationalGaussian.predict` runs `add_jitter`
on its output covariance, so the per-point marginals `elbo` sees carry the
same offset, and without it `dual_elbo` would miss `elbo` at matched moments
by N*eps/(2 sigma^2).

`dual_elbo` is the same functional as `elbo` but evaluated as a function of
the sites and the hyperparameters. Its value matches `elbo` at the implied
moments (measured 5.7e-14 absolute, 2.8e-16 relative at random PSD sites)
while its hyperparameter gradient differs, because q moves with theta through
Kzz while the sites stay frozen. Kzz is deliberately not detached and no
moments are cached on the module; a cached implementation would pass every
value assertion and fail only `test_dual_elbo_hyper_gradients_differ_away_
from_optimum`.

`natural_gradient_step` and `variational_coordinates` gain dual
registrations. Because grad_mu KL = lambda exactly, the KL is never
differentiated and the step is a convex combination towards a closed-form
target built from one `jax.grad` of `expected_log_likelihood` (Bonnet and
Price), with N/B scaling, a trace-safe floor on beta and a symmetrise. That
makes it the Salimbeni step at gamma = rho: matched initialisations agree to
1.4e-15 in (m, S) over six Bernoulli steps at every rate tested. One rho = 1
full-batch conjugate step lands on the Titsias optimum to 4.9e-15.

`fit_natgrads` rejects a numeric step size above one on this family, since
the update is a convex combination; schedules cannot be checked statically
and are documented as the caller's responsibility. Plain `fit` on the family
also works and is documented as gradient descent in the dual coordinates.

The `VariationalParametrisationSuite` benchmark gains a `dual` arm, in both
`benchmarks/objectives.py` and the hardcoded parametrize list in
`tests/test_benchmarks_smoke.py`.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
Adds examples/natgrads.py, the first of two tutorial notebooks for the
natural-gradient stack. It derives the exponential-family view of q(u),
shows that the Fisher information is the Jacobian d(eta)/d(theta) (checked
numerically to 2.9e-15), reads the update as mirror descent, and then runs
two demos:

* a conjugate 1D regression where one gamma=1 full-batch natural-gradient
  step recovers the Titsias optimum to 1.2e-13 while Adam on the same
  problem is still 4.6 nats short after 2000 iterations;
* a mini-batched 2D banana Bernoulli benchmark (N=2000, M=50, B=256, 1000
  iterations) comparing natural gradients + Adam against Adam alone, per
  iteration and per wall-clock second.

Closes with the negative-definite cone result, a gamma sweep reproducing
its boundary, and a demonstration of the step-size backoff.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
`mkdocs build` aborted with `IndexError: string index out of range` while
rendering `_examples/dual_svgp.md`. The trigger is a markdown-katex parser
bug: `iter_inline_katex` reads `line[end + 1]` without a bounds check, so any
line whose final characters are a backtick code span immediately preceded by
`$` crashes the build. The prose read

    ... that is $-$`sparsity_gap`
    above, and ...

where the closing `$` of `$-$` abuts the code span and the span ends the line.

Reword to "the negated `sparsity_gap` computed above", which removes the
`$`-adjacent code span entirely rather than relying on a particular line wrap.
The meaning is unchanged, and the sentence still says c(theta) is minus the
Titsias trace term. A scan of all generated `docs/_examples/*.md` confirms this
was the only occurrence of the pattern in the docs tree.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
CHANGELOG: add the missing "### Added" entry for fit_natgrads and
gpjax.natural_gradients. PR#1 shipped both without a changelog entry, so the
Removed entry added here forward-referenced an API the changelog never
announced. Bringing it forward from PR#3 keeps any release cut mid-stack
self-consistent; PR#3 appends the dual entries to the same section.

CHANGELOG: correct the justification prose. "the natural gradient with respect
to theta equals the ordinary gradient with respect to eta, in any
parameterisation" is false as literally written -- for a reparameterisation xi
with J = dtheta/dxi, the natural gradient in xi is J^-1 grad_eta L, not
grad_eta L. The identity is specific to the natural/expectation pair of an
exponential family. Reworded to state that pairing and the point it supports:
either coordinate system is recoverable on the fly, so no dedicated class is
needed.

benchmarks: drop diff-relative wording from the
VariationalParametrisationSuite docstring. "surviving" only means something to
someone reading this commit's diff, and "Both" is a count PR#3 invalidates when
it re-adds the dual arm. The module docstring keeps its explicit
"(standard, whitened)" list, which PR#3 must extend regardless.

tests: split the _psd guard out of test_removed_families_are_gone. The _psd arm
asserted `"_psd" not in __all__`, vacuous for a helper that was never exported,
under a failure message about superseded parameterisations. It is now its own
test with a docstring saying what it actually guards.

CLAUDE.md: "Three optimisers" -> four. fit_natgrads landed in gpjax/fit.py in
PR#1; the sentence has been stale since.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
Factor R in the Kzz basis. Forming R = Kzz + Kzz Lambda_2 Kzz explicitly
carries a rounding error of order ||Kzz||^2 ||Lambda_2|| eps, which for a
large-variance kernel (or in float32) exceeds lambda_min(R) ~= jitter, so
chol(R) returned NaN and poisoned every later fit_natgrads iterate --
measured at RBF variance 1e4, M=80, default jitter, float64, where the
matched VariationalGaussian run stayed finite. R = Lk (I + Lk^T Lambda_2 Lk)
Lk^T is factorised instead, giving a lower-triangular Lr = Lk chol(I + G)
whose inner matrix has lambda_min >= 1 - O(||G|| eps). Same two Choleskys
per call, plus one M x M product. The "chol(R) never fails" and "no Cholesky
a backoff could rescue" claims are softened to match.

Also: split _gram_and_root off _working_matrices so the dual natgrad step
stops discarding a chol(R); take tr(R^-1 Kzz) as ||Lr^-1 Lk||_F^2 rather
than a full cho_solve; bound-check an optax schedule against rho <= 1 over
the whole num_iters horizon for the dual family, which previously returned
a silent all-NaN history; hoist the _fmt_Kzt_Ktt/_fmt_inducing_inputs hooks
to AbstractVariationalGaussian and keep one typed _symmetrise; convert the
dual family's numpydoc sections to the Google style the file and mkdocs use,
and document marginals' inputs argument.

Doc corrections, all measured: elbo on a DualVariationalGaussian returns the
same value and the same gradients as dual_elbo (bit-identical value, 1.7e-14
on gradients), so the CHANGELOG's gradient claim now names the matched
VariationalGaussian as the comparison; and vmap does not repeat the
unbatched factorisations per datum, so neither elbo nor the benchmark arm
pays 2N Choleskys -- eager counts are 4 (dual), 2 (standard), 1 (whitened),
and the compiled dual_elbo and dual step are 2 potrf each.

Tests: regression for chol(R) at variance 1e4/M=80 and 1e3/M=50 with the
default jitter, an Lr Lr^T = R reconstruction check, the schedule guard, a
half-batch arm on the dual/elbo equivalence plus a direct pin on the N/B
factor, and the triplicated dual fixtures moved to tests/_dual_helpers.py.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
Correct the mini-batch ramp argument (the target is q-dependent outside
conjugacy, so gamma=1 lands on a moving fixed-point target, not the
mini-batch optimum), replace the unmeasured calibration claim after the
banana contours with the metrics the cell actually prints, and attribute
the cone sweep's gamma=2 failure to the over-confident S_0 rather than to
gamma=2 itself.

Smaller corrections: the Fisher solve is O((M + M(M+1)/2)^3) in the vec_s
coordinates, not O((M + M^2)^3); the one-step demo agrees to ~1e-13, not
fourteen decimal places; the sparse/exact predictive deviation is
quantified and located outside the data range; the Adam ELBO-gap
description now matches the shape of the log-log panel; the roadmap says
"exact variational optimum" where the notebook later reserves "exact
posterior" for the non-sparse GP; K=100 is explained against the paper's
dataset-dependent K.

Code: derive the ELBO figure's y-limits from the smoothed histories so
neither curve is clipped, guard the crossing report against never
reaching the target, draw the held-out points in the categorical palette
with per-class markers instead of the contour colourmap, use banana_data
in the split report, and run ruff format (the pinned make_banana block is
byte-identical afterwards).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
The dual/Salimbeni E-step divergence on the banana demo was attributed to
floating-point conditioning. It is not: `inv_probit` clips its output into
[1e-3, 1-1e-3], so the computed Bernoulli log-likelihood is not log-concave in
the tails (positive second derivative for f < -2.44). A confidently mislabelled
point then yields beta_i < 0, the dual branch's beta_floor clips it, and the two
branches diverge. With the clip disabled the same six steps agree to 6.2e-13
instead of 5.1e-3.

- Re-attribute the mechanism in the dual notebook and add a diagnostic cell that
  measures it, and qualify the "identical iterates" claim wherever it is stated
  (natural_gradients.py, fit.py, CHANGELOG, both notebooks): it holds provided
  the computed beta stays non-negative.
- Note in the natgrads notebook that the cone discussion assumes a log-concave
  *computed* likelihood, which the clipped probit violates in the far tails.
- Reword the Salimbeni registration docstrings: dispatch covers
  GraphVariationalGaussian, but the standard elbo path is broken upstream
  (MatrixLinearOperator dimensionality error out of gram), on main too.
- Extend _check_natgrad_schedule to reject non-positive rates for every family,
  not just the dual upper bound.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
The Sphinx docs arrived with v1.0, after this branch was cut, so the
deletion of NaturalVariationalGaussian and ExpectationVariationalGaussian
now has to reach three doc files the original commits could not know
about: drop the two classes from the variational-families autosummary
page, and repoint the glossary's "natural parameters" entry from the
removed classes to fit_natgrads on the surviving families. fit_natgrads
is added to the fit reference page so that glossary link resolves —
an omission from PR#1, which shipped the function without a reference
entry.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
…e v1.0 API

The v1.0 likelihoods are pure conditionals, so the minibatch ELBO scale in
dual_elbo and the dual natural-gradient step now derives from
Dataset.n_total instead of likelihood.num_datapoints, the constructor
annotation follows the AbstractPosterior -> JointModel split, and the
docstring examples plus the dual test fixtures drop num_datapoints. The
minibatch tests stamp n_total onto their batch views to preserve the N/B
factor they assert.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
thomaspinder and others added 9 commits August 7, 2026 22:58
…tation

AbstractKernelComputation.gram() called self.cross_covariance(kernel, x, x)
directly instead of self._gram(kernel, x), leaving _gram (and any subclass
override of it) dead code. This silently broke BasisFunctionComputation's
_gram override, which computes the RFF feature matrix once: every
RFF.gram(x) call was routing through cross_covariance's _cross_covariance
override instead, which independently computes the feature matrix twice
(once as z1, once as z2) even though x is the same array both times.

gram() now calls self._gram(kernel, x), with the base class's _gram
implemented as cross_covariance(kernel, x, x) -- identical behaviour to
before for every kernel that doesn't override _gram (RBF, Matern, etc. via
DenseKernelComputation). DiagonalKernelComputation and
ConstantDiagonalKernelComputation override gram() itself, not _gram, so
they are unaffected. RFF now correctly uses its single-pass _gram override,
computing the feature matrix once instead of twice.

Adds a monkeypatched call-count test that fails before this fix (2 calls)
and passes after (1 call), plus numerical-equivalence tests confirming
gram() output is unchanged for RBF, Matern12/32/52, and RFF kernels.

Fixes #679.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
…init workaround, fix White's phantom lengthscale

Scoped fix for issue #695 per plans/2026-07-26-issue-695-abstract-final.md
(the "measured middle path", not a wholesale abstract/final rewrite).

1. Give StationaryKernel/AbstractKernel fields real defaults matching their
   __init__ defaults (lengthscale/variance via eqx.field(default_factory=...),
   compute_engine via DenseKernelComputation). Equinox inherits hand-written
   __init__ methods at runtime, but pyright's @dataclass_transform synthesises
   a fresh __init__ per subclass from declared fields whenever a class (e.g.
   RBF, Matern12/32/52) doesn't define its own -- so undefaulted fields made
   the canonical `RBF()` a pyright error while `RBF(name="xyz")` type-checked
   cleanly (TypeError at runtime). `name` is now a ClassVar on every concrete
   stationary kernel, matching runtime (it was never a real __init__ param)
   and removing it from pyright's synthesised signature.

   Note: ClassVar must be imported from stdlib `typing`, not
   `beartype.typing` -- despite `beartype.typing.ClassVar is typing.ClassVar`
   at runtime, pyright's dataclass field-exclusion does not recognise the
   re-exported symbol and keeps such fields in the synthesised __init__.
   Verified against pyright 1.1.411.

2. Removed `_compute_base_init` (kernels/base.py). Its docstring claimed
   equinox modules are frozen once any parent __init__ returns, which no
   longer holds under equinox 0.13.8 (verified directly: a child class can
   set its own fields after `super().__init__()` returns). StationaryKernel
   now calls plain `super().__init__(...)`.

3. White hardcoded lengthscale=1.0 into StationaryKernel.__init__ even
   though White.__call__ (and its inherited spectral_density, which raises
   before touching lengthscale) never reads it, so every White carried a
   real, trainable PositiveReal leaf with zero gradient that polluted
   optimiser/MCMC state. White now overrides `lengthscale` as a ClassVar
   (`None`) rather than inheriting it as a dataclass field, and calls
   `AbstractKernel.__init__` directly instead of `StationaryKernel.__init__`,
   so lengthscale is absent from White's pytree entirely (1 leaf instead of
   2 -- verified via jax.tree_util.tree_flatten).

The Zero mean-function fix from the same research note was already in place
(gpjax/mean_functions.py) and is unaffected by this change; reverified via
the note's repro script.

Residual, out of scope: SumKernel/ProductKernel (kernels/base.py) still
synthesise a `kernels: tuple` parameter under pyright that rejects a
`list[AbstractKernel]` argument (`SumKernel(kernels=[...])`). This was
already broken pre-fix (masked by RBF()'s own missing-argument errors) and
sits one level deeper -- fixing it needs a converter on the `kernels` field,
which is a different kind of change to the one scoped here.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
…683)

Both functions implemented the identical power-sum/sign/mask/fori_loop
Newton-Girard recursion, differing only in the trailing broadcast shape
of the per-dimension values (scalar in oak.py vs. (N, N) matrix in
sobol.py). Generalised gpjax/kernels/additive/oak.py::_newton_girard to
accept an arbitrary trailing shape via broadcasting (exponents/signs/mask
reshaped against z.ndim, elem_sym shaped from z.shape[1:]), and made
sobol.py import and call it directly, deleting _newton_girard_matrices.

Pure refactor: replaced the scalar-only jnp.dot reduction with
jnp.sum(product, axis=0), which computes the identical sum and matches
what the matrix variant already did, so the scalar call site's output is
numerically unchanged. Verified via existing OAK/Sobol tests (unchanged,
all passing) plus a new test asserting the matrix-shaped path reduces to
the same per-element scalar recursion.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
The soak-period precondition from #637's rollout plan (benchmarks/ and
asv-constraints.txt landing on main) has been satisfied for months, so
per plans/2026-05-01-continuous-benchmarking-impl.md Task 12, bench-check
now runs as part of `uv run poe all-tests`, between docstrings and test.

Updates benchmarks/README.md's "Notes on bench-check" section, which
previously explained why the task wasn't yet wired in, to describe the
current wiring instead. Adds tests/test_poe_tasks.py asserting the
all-tests sequence and bench-check task definition via tomllib, following
the existing pyproject.toml-parsing pattern in tests/test_dependencies.py.

Closes #638.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
…kelihood/covariance API

The fix commit was cherry-picked from a worktree that was accidentally
based on main instead of v1.0; its own new test still used the
pre-v1.0 Gaussian(num_datapoints=...) constructor and the
return_covariance_type= kwarg, both removed by the v1.0 rename sweep.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
Closes #515. Adds a way to regularise gradient-based MLE fitting with
hyperparameter priors (GPyTorch-style MAP), e.g. strongly preferring
large lengthscale/noise over overfitting-prone small ones in high
dimensions.

Design: `with_log_prior(objective, log_prior)` is a plain combinator
over the existing `Objective = Callable[[Model, Dataset], ScalarFloat]`
protocol -- it returns a new objective computing
`objective(model, data) + log_prior(model)`, where `log_prior` is a
user-supplied function over the model pytree. This composes directly
with `fit`/`fit_scipy`/`fit_lbfgs` with no changes to those functions,
and conjugate_mll/other objectives are untouched when the feature is
unused.

This deliberately does NOT reintroduce the pre-v1 `Parameter(...,
prior=...)` field removed in 98a7feb (#621). That mechanism attached a
prior to every Parameter class, which (a) tangled the constrained/
unconstrained bijection with an ambiguous "is this prior for gradient
descent or for NumPyro?" scope (the numpyro-priors.md design doc that
preceded the removal explicitly called this out and worked around it
with a separate `numpyro_properties` namespace), and (b) duplicated the
fully Bayesian path, which now samples hyperparameters directly via
`numpyro.sample` and feeds raw arrays into GPJax constructors (see
`tests/test_numpyro_extras.py`). A standalone objective combinator
avoids both problems: no field on Parameter, no coupling to NumPyro,
and the fully-Bayesian path is untouched.

Known simplification, documented in the docstring: `log_prior` is
evaluated on the constrained (unwrapped) parameter values -- the same
values `objective` receives -- without the change-of-variables Jacobian
for the unconstrained space the optimiser actually moves in. For the
strongly regularising priors this feature targets, that's an
acceptable simplification rather than a literal Bayesian MAP.

Tests added to tests/test_objectives.py:
- test_with_log_prior_matches_base_objective_when_prior_is_zero: a
  zero log-prior leaves both value and gradient identical to the
  unregularised objective.
- test_with_log_prior_map_regularised_fit_prefers_prior_consistent_lengthscale:
  end-to-end MAP demonstration -- an unregularised fit_scipy fit on a
  high-frequency, low-noise signal collapses the RBF lengthscale to a
  tiny, overfitting value, while wrapping the same objective with a
  LogNormal(log(3), 0.15) prior on the lengthscale converges near the
  prior mean instead.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
Adds a heavy-tailed Student's t observation likelihood (Jylanki,
Vanhatalo & Vehtari 2011) following the Bernoulli/Poisson pattern:
GHQuadratureIntegrator(20) for the expected log-likelihood, since the
Student's t is not conjugate to a Gaussian latent. degrees_of_freedom
and scale are both PositiveReal-wrapped.

Ported from a fix implemented against a stale pre-v1.0 checkout (the
original used the removed Gaussian-style num_datapoints constructor
argument); adapted here to v1.0's pure-conditional likelihood
contract and NonConjugateModel.init_latent() lazy-latent pattern.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
@github-actions

github-actions Bot commented Aug 7, 2026

Copy link
Copy Markdown

📖 Docs preview: https://pr-752--endearing-crepe-c2d5fe.netlify.app

Smoke render — the expensive notebooks run with reduced budgets, so
figures are not publication fidelity. /render-mode.txt says smoke.

thomaspinder and others added 20 commits August 8, 2026 06:04
Fix bug when unwrapping paramax parameters
# Conflicts:
#	gpjax/state_space/inference.py
#	tests/test_state_space/test_prediction.py
…ovariance

feat(state-space): dense joint predictive covariance
Fold the CLAUDE.md guidance into AGENTS.md, which becomes the single
tracked agent-guidance file, and drop it from .gitignore.

Content that moved across was checked against the code first. Corrections
made while merging:

- Posterior -> JointModel. `Prior * Likelihood` dispatches through
  construct_model() to ConjugateModel / NonConjugateModel; Posterior
  (ExactPosterior / LatentPosterior) is what model.condition(data)
  returns. Heteroscedastic models cannot be built with `*`.
- gpjax/numpyro_extras.py no longer exists; describe the numpyro
  constraint/bijector use and the docs/examples/numpyro_integration.py
  MCMC pattern instead.
- all-tests is lint + docstrings + bench-check + test.
- conftest lives in tests/ and installs the jaxtyping import hook with
  beartype; `poe test` skips tests marked slow.

Stale claims already in AGENTS.md were corrected too: ruff replaces
Black/isort, the docs are Sphinx + MyST-NB under docs/examples rather
than MkDocs under examples/, and `poe docs-build` is now `poe docs`.

Also track .github/workflows/zizmor.yml.
zizmor.yml was added but never functional: `uv run --frozen --only-group
dev zizmor` had no zizmor to run, and the `poe lint-actions` task named
in its error message did not exist. Add both, then fix everything the
audit reported (96 findings -> 0).

Tooling:
- zizmor in the dev group, resolved through uv.lock like ruff, so a new
  release cannot turn every open PR red without a commit.
- `poe lint-actions` runs the annotated form locally. Unlike the SARIF
  the workflow renders, plain output exits 14 on findings, so the task
  gates on its own exit code.

Findings:
- unpinned-uses (61): every action pinned to a commit SHA with the
  release tag in a trailing comment. Dependabot bumps these.
- artipacked (13): persist-credentials: false on every checkout. No job
  pushes with git; the release job cuts its release through the API.
- excessive-permissions (9): workflow-level scopes moved down to the one
  job that needs them, with `permissions: {}` at the top. The jobs that
  execute repository code — tests, scans, notebook renders — now run
  with a read-only token, and only create-release writes contents, only
  deploy-docs and publish-pypi hold a write or OIDC scope.
- template-injection (3): release.yml interpolated the workflow_dispatch
  version input and two output values straight into run blocks; they go
  through env now.
- cache-poisoning (6): release.yml no longer caches at all, and the
  Sunday full render disables its caches on release-triggered runs.
  zizmor cannot see through the guard expression, so those three carry a
  scoped ignore with the reasoning.
- dangerous-triggers (1): auto-label needs pull_request_target to label
  fork PRs. Its checkout step was unnecessary — actions/labeler reads
  its config through the API — so the workflow no longer fetches PR code
  at all, and the trigger carries a scoped ignore.
- superfluous-actions (1): softprops/action-gh-release replaced with
  `gh release create`, removing a third-party action from the publishing
  path.
- dependabot-cooldown (2): seven-day cooldown on both update entries.

Fixes a latent bug found while editing auto-label.yml: three steps were
guarded on `github.event_name == 'pull_request'`, which is never true
under `pull_request_target`, so PR labelling, size labelling and
content labelling had all been dead.

Verification: `uv run poe lint-actions` reports no findings, and the
workflow's own gate (`jq '[.runs[].results[]] | length'` over the SARIF)
counts 0.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants